# -*- coding: utf-8 -*-
# RSA public and private key class
# See http://en.wikipedia.org/wiki/RSA_(algorithm)
#     and http://www.di-mgt.com.au/rsa_alg.html
"""RSA public and private key module

NOTICE: This program is used for learning, there may have bugs
and potential security issues. Please DO NOT use this program
in production environment.
"""

__author__ = 'Kainan Zhu'
__email__ = 'dlangu0393@gmail.com'
__copyright__ = 'Copyright 2013, Kainan Zhu'

from random import randrange, getrandbits
from itertools import repeat
from rsa.util import bit_length


class RSAKey(object):
    """Base class for all RSA key class"""


class PublicKey(RSAKey):
    """
    RSA Public key class

    The syntax in PKCS#1:
        RSAPublicKey ::= SEQUENCE {
            modulus         INTEGER,    -- n
            publicExponent  INTEGER     -- e
        }
    """

    __slots__ = ('n', 'e', 'tuple')

    def __init__(self, n, e):
        self.n = n
        self.e = e

    def __eq__(self, other):
        if other is None:
            return False

        if not isinstance(other, PublicKey):
            return False

        return self.n == other.n and self.e == other.e

    def __ne__(self, other):
        return not self == other

    def __repr__(self):
        return 'PublicKey(%d, %d)' % self.tuple()

    def tuple(self):
        """
        Get the public key as (e, n) tuple.
        :return:
            A tuple (e, n).
        """
        return self.e, self.n


class PrivateKey(RSAKey):
    """
    RSA Private key class

    The syntax in PKCS#1:
        RSAPrivateKey ::= SEQUENCE {
            version             Version,
            modulus             INTEGER,    -- n
            publicExponent      INTEGER,    -- e
            privateExponent     INTEGER,    -- d
            prime1              INTEGER,    -- p
            prime2              INTEGER,    -- q
            exponent1           INTEGER,    -- d mod (p-1)
            exponent2           INTEGER,    -- d mod (q-1)
            coefficient         INTEGER,    -- (inverse of q) mod p
            otherPrimeInfos     OtherPrimeInfos OPTIONAL
        }
    """

    __slots__ = ('n', 'e', 'd', 'p', 'q', 'exp1', 'exp2', 'coe', 'tuple')

    def __init__(self, n, e, d, p, q, exp1=None, exp2=None, coe=None):
        self.n = n
        self.e = e
        self.d = d
        self.p = p
        self.q = q

        if exp1 is None:
            self.exp1 = int(d % (p - 1))
        else:
            self.exp1 = exp1

        if exp2 is None:
            self.exp2 = int(d % (q - 1))
        else:
            self.exp2 = exp2

        if coe is None:
            # c = q^-1 mod p
            # q^-1 ≡ c (mod p)
            self.coe = mod_inverse(q, p)
        else:
            self.coe = coe

    def __eq__(self, other):
        if other is None:
            return False

        if not isinstance(other, PrivateKey):
            return False

        return (self.n == other.n
                and self.e == other.e
                and self.d == other.d
                and self.p == other.p
                and self.q == other.q
                and self.exp1 == other.exp1
                and self.exp2 == other.exp2
                and self.coe == other.coe)

    def __ne__(self, other):
        return not self == other

    def __repr__(self):
        return 'PrivateKey(%d, %d, %d, %d, %d, %d, %d, %d)' \
               % (self.n, self.e, self.d, self.p, self.q, self.exp1, self.exp2, self.coe)

    def tuple(self):
        """
        Get the private key as (n, e, d, p, q) tuple.
        :return:
            A tuple (n, e, d, p, q).
        """
        return self.n, self.e, self.d, self.p, self.q


def generate_prime(nbits):
    """
    Generate a tuple of two random primes of length bitLength.
    :param nbits:
        Integer value. Indicating the length of primers generated in bits.
    :return:
        A tuple (p, q) in which p, q are n-bit, in which n is specified by bitLength param, random primes.
    """

    def is_prime(p, t=7):
        """
        Check if p is a prime using Miller-Rabin primality test
        See http://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test

        The Miller-Rabin primality test in pseudo-code:
            Input: n > 3, an odd integer to be tested for primality;
            Input: k, a parameter that determines the accuracy of the test
            Output: composite if n is composite, otherwise probably prime
            write n − 1 as 2s·d with d odd by factoring powers of 2 from n − 1
            WitnessLoop: repeat k times:
               pick a random integer a in the range [2, n − 2]
               x ← ad mod n
               if x = 1 or x = n − 1 then do next WitnessLoop
               repeat s − 1 times:
                  x ← x2 mod n
                  if x = 1 then return composite
                  if x = n − 1 then do next WitnessLoop
               return composite
            return probably prime

        :param t:
            A parameter that determines the accuracy of the test.
        :param p:
            The number to be tested.
        :return:
            A boolean value.
        """

        def is_composite(n, s, d):
            """
            Check if n is composite
            """
            if pow(n, d, p) == 1:
                return False
            for i in range(s):
                if pow(n, 2 ** i * d, p) == p - 1:
                    return False
            return True

        # Dealing the case where n is less than 3
        if p < 3:
            return [False, False, True][p]
        elif not p & 1:
            return False
        else:
            s, d = 0, p - 1
            while not d & 1:
                s += 1
                d >>= 1
        for _ in repeat(None, t):  # Use repeat function in itertools for efficient looping
            if is_composite(randrange(2, p), s, d):
                return False
        return True

    p = getrandbits(nbits)
    while not is_prime(p):
        p = getrandbits(nbits)
    return p


def generate_p_q(bitLength):
    """
    Generate a pair of p and q which will not be equal and will be different enough
    :param bitLength:
        The bit length of p and q.
    :return:
        p, q
    """

    total_length = bitLength * 2

    # This expression was found in sybren's python-rsa module
    # Reference was not found
    offset = bitLength // 16

    p_length = bitLength + offset
    q_length = bitLength - offset

    p, q = 0, 0
    while p == q or bit_length(p * q) != total_length:
        p, q = generate_prime(p_length), generate_prime(q_length)

    return p, q


def egcd(a, b):
    """
    Extended Euclidean Algorithm
    Returns x, y, gcd(a,b) such that ax + by = gcd(a,b).
    See http://en.wikipedia.org/wiki/Extended_Euclidean_algorithm
        and http://en.wikibooks.org/wiki/Algorithm_Implementation/Mathematics/Extended_Euclidean_algorithm

    Extended Euclidean Algorithm in pseudo-code:
        function extended_gcd(a, b)
            x := 0    lastx := 1
            y := 1    lasty := 0
            while b ≠ 0
                quotient := a div b
                (a, b) := (b, a mod b)
                (x, lastx) := (lastx - quotient*x, x)
                (y, lasty) := (lasty - quotient*y, y)
            return (lastx, lasty)

    :param a:
    :param b:
    :return:
    """
    x, y, u, v = 0, 1, 1, 0
    while a != 0:
        q, r = b // a, b % a
        m, n = x - u * q, y - v * q
        b, a, x, y, u, v = a, r, u, v, m, n
    return b, x, y


def mod_inverse(a, m):
    """
    Modular multiplicative inverse
    See http://en.wikibooks.org/wiki/Algorithm_Implementation/Mathematics/Extended_Euclidean_algorithm

    To get d such that de ≡ 1 (mod m):
    the modular multiplicative inverse of an integer a modulo m is an integer x such that
        a^-1 ≡ x (mod m)
    which is equal to
        a ≡ x^-1 (mod m)
        ax ≡ 1 (mod m)
        ax - 1 = cm
    finally, we can get
        ax - cm = 1

    the egcd(a, m) returns gcd(a, m), x, y such that
        ax + ym = gcd(a, m)
    since a and m are coprime
        gcd(a, m) = 1
    hence
        ax + ym = 1
    then we can find that x returned by egcd(a, m) is exactly what we want

    Then we need 1 < d < m:
        x % m

    :param a:
    :param m:
        a must be coprime to m.
    :return:
        x such that a^-1 ≡ x (mod m).
    """
    divisor, mod_inv, _ = egcd(a, m)

    if not divisor == 1:
        raise ValueError('Param a and m are not coprime')

    return mod_inv % m


def generate_keys(p, q):
    """
    Generate RSA public and private keys.
    :param p:
    :param q:
        A tuple (pub, pri) in which are instances of PublicKey and PrivateKey.
    """
    phi, n = (p - 1) * (q - 1), p * q
    # 65537 is a commonly used value for e
    # See http://en.wikipedia.org/wiki/RSA_(algorithm)#Faulty_key_generation
    return PublicKey(n, 65537), PrivateKey(n, 65537, mod_inverse(65537, phi), p, q)


def new_keys(nbits):
    """
    Generates public and private keys
    :param nbits:
        The key length.
    :return:
        The public and private keys as ((e,n), (d, n)).
    """
    p, q = generate_p_q(nbits // 2)
    return generate_keys(p, q)


__all__ = ['new_keys', 'generate_keys', 'generate_p_q', 'generate_prime',
           'RSAKey', 'PublicKey', 'PrivateKey']